#pragma once

#include <utility>
#include <iostream>
#include <assert.h>
#include <algorithm>
using namespace std;

template<class K, class V>
struct AVLTreeNode
{
	typedef AVLTreeNode<K, V> Node;
	pair<K, V> _kv;
	Node* _left;
	Node* _right;
	Node* _parent;
	int _bf;
	AVLTreeNode(const pair<K, V>& kv) :_left(nullptr), _right(nullptr),_parent(nullptr), _kv(kv), _bf(0) { }
};

template<class K, class V>
class AVLTree
{
private:
	typedef AVLTreeNode<K, V> Node;
	Node* _root = nullptr;
	int _size = 0;
public:
	bool Insert(const pair<K, V>& kv)
	{
		if (_root == nullptr)
		{
			_root = new Node(kv);
			_size++;
			return 1;
		}

		Node* parent = nullptr, * cur = _root;
		while (cur)
		{
			if (kv.first == cur->_kv.first) return 0;
			parent = cur;
			if (kv.first > cur->_kv.first) cur = cur->_right;
			else cur = cur->_left;
		}
		cur = new Node(kv);
		if (parent->_kv.first > kv.first)
		{
			parent->_left = cur;
			parent->_left->_parent = parent;
		}
		else
		{
			parent->_right = cur;
			parent->_right->_parent = parent;
		}


		while (parent)
		{
			if (parent->_left == cur) parent->_bf--;
			else parent->_bf++;
			if (parent->_bf == 0) break;
			else if (parent->_bf == 1|| parent->_bf == -1)
			{
				cur = parent;
				parent = parent->_parent;
			}
			else if (parent->_bf == 2 || parent->_bf == -2)
			{
				if (parent->_bf == -2 && parent->_left->_bf == -1)
				{
					RotateL(parent);
				}
				else if (parent->_bf == 2 && parent->_right->_bf == 1)
				{
					RotateR(parent);
				}
				else if (parent->_bf == -2 && parent->_left->_bf == 1)
				{
					RotateLR(parent);
				}
				else if (parent->_bf == 2 && parent->_right->_bf == -1)
				{
					RotateRL(parent);
				}
				break;
			}
			else
			{
				assert(0);
			}
		}
		_size++;
		return 1;
	}

	void Erase(const K& k)
	{
		Node* parent = nullptr, * cur = _root;
		while (cur)
		{
			if (k > cur->_kv.first)
			{
				parent = cur;
				cur = cur->_right;
			}
			else if (k < cur->_kv.first)
			{
				parent = cur;
				cur = cur->_left;
			}
			else break;
		}
		if (cur == nullptr) return;
		if (cur->_right == nullptr)
		{
			Node* child = cur->_left;
			if (parent == nullptr)
			{
				_root = child;
				if (child != nullptr)
					_root->_parent = nullptr;
				delete cur;
			}
			else 
			{
				if (parent->_left == cur)
				{
					parent->_left = child;
					if(child != nullptr)
						child->_parent = parent;
					parent->_bf++;
				}
				else
				{
					parent->_right = child;
					if (child != nullptr)
						child->_parent = parent;
					parent->_bf--;
				}
				delete cur;
			}
		}
		else if (cur->_left == nullptr)
		{
			Node* child = cur->_right;
			if (parent == nullptr)
			{
				_root = child;
				_root->_parent = nullptr;
				delete cur;
			}
			else
			{
				if (cur == parent->_left)
				{
					parent->_left = child;
					child->_parent = parent;
					parent->_bf++;
				}
				else
				{
					parent->_right = child;
					child->_parent = parent;
					parent->_bf--;
				}
				delete cur;
			}
		}
		else
		{
			Node* RightMin = cur->_right, * RightMinParent = cur;
			while (RightMin->_left)
			{
				RightMinParent = RightMin;
				RightMin = RightMin->_left;
			}
			swap(cur->_kv, RightMin->_kv);
			if (RightMinParent == cur) cur->_right = nullptr;
			else RightMinParent->_left = nullptr;
			delete RightMin;
		}
	}

	int size()
	{
		return _size;
	}

	int height()
	{
		return _Height(_root);
	}

	V& operator[](const K& k)
	{
		Node* cur = _root;
		Insert({k, V()});
		while (cur)
		{
			if (cur->_kv.first > k)
			{
				cur = cur->_left;
			}
			else if (cur->_kv.first < k)
			{
				cur = cur->_right;
			}
			else return cur->_kv.second;
		}
	}

	bool Empty()
	{
		return _root == nullptr;
	}

	void InOrder()
	{
		_InOrder(_root);
		cout << "------------------------------------------------" << endl;
	}

	bool IsAVTree()
	{
		return _IsAVTree(_root);
	}

private:
	int _Height(Node* root)
	{
		if (!root) return 0;
		return max(_Height(root->_left), _Height(root->_right)) + 1;
	}

	bool _IsAVTree(Node* root)
	{
		if (!root) return 1;
		return _IsAVTree(root->_left) && _IsAVTree(root->_right)
			&& abs(_Height(root->_left) - _Height(root->_right)) < 2;
	}

	void _InOrder(Node* root)
	{
		if (root == nullptr) return;
		_InOrder(root->_left);
		cout << root->_kv.first << ':' << root->_kv.second << endl;
		_InOrder(root->_right);
	}

	void RotateLR(Node* root)
	{
		Node* SubL = root->_left;
		Node* SubLR = SubL->_right;
		int bf = SubLR->_bf;
		RotateR(SubL);
		RotateL(root);
		root->_bf = SubL->_bf = SubLR->_bf = 0;
		if (bf == -1)
			root->_bf = 1;
		else if (bf == 1)
			SubL->_bf = -1;
	}

	void RotateRL(Node* root)
	{
		Node* SubR = root->_right;
		Node* SubRL = SubR->_left;
		int bf = SubRL->_bf;
		RotateL(SubR);
		RotateR(root);
		root->_bf = SubR->_bf = SubRL->_bf = 0;
		if (bf == -1)
			SubR->_bf = 1;
		else if (bf == 1)
			root->_bf = -1;
	}

	void RotateR(Node* root)
	{
		Node* SubR = root->_right;
		Node* SubRL = SubR->_left;
		Node* Parent = root->_parent;

		if (SubRL != nullptr)
			SubRL->_parent = root;
		root->_right = SubRL;
		root->_parent = SubR;
		SubR->_left = root;

		if (Parent == nullptr)
		{
			_root = SubR;
			SubR->_parent = nullptr;
		}
		else
		{
			SubR->_parent = Parent;
			if (root == Parent->_left) Parent->_left = SubR;
			else Parent->_right = SubR;
		}
		root->_bf = 0;
		SubR->_bf = 0;
	}

	void RotateL(Node* root)
	{
		Node* SubL = root->_left;
		Node* SubLR = SubL->_right;
		Node* Parent = root->_parent;

		if (SubLR != nullptr)
			SubLR->_parent = root;
		root->_left = SubLR;
		root->_parent = SubL;
		SubL->_right = root;

		if (Parent == nullptr)
		{
			_root = SubL;
			_root->_parent = nullptr;
		}
		else
		{
			SubL->_parent = Parent;
			if (Parent->_left == root) Parent->_left = SubL;
			else Parent->_right = SubL;
		}
		root->_bf = 0;
		SubL->_bf = 0;
	}
};